# Author : ZZH
# Date : 2025/2/6
"""
生成视频摘要
"""
from transformers import pipeline
import re

# 加载文本生成管道
generator = pipeline("text-generation", model="gpt2-medium")  # 使用中等规模模型


def clean_summary(text):
    """清理生成文本中的重复内容"""
    text = re.sub(r"\n+", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


# 读取帧描述
with open("frame_descriptions.txt", "r") as f:
    context = " ".join([line.strip() for line in f.readlines()])

# 构建提示词
prompt = f"""Based on the following observations from video frames:
{context}

Generate a concise video summary in 3 sentences:"""

# 生成摘要
summary = generator(
    prompt,
    max_length=200,
    num_return_sequences=1,  # 返回的序列数量
    temperature=0.7,         # 控制创造性（值越小越保守）
    do_sample=True           # 是否使用采样方法生成文本
)[0]['generated_text']

# 提取有效部分
summary = summary.split("Generate a concise video summary in 3 sentences:")[-1]
print("Video Summary:\n", clean_summary(summary))
